package io.summer.gateway.filter;

import io.summer.gateway.tool.MediaTypeTool;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.cloud.gateway.handler.AsyncPredicate;
import org.springframework.cloud.gateway.handler.predicate.AbstractRoutePredicateFactory;
import org.springframework.cloud.gateway.handler.predicate.ReadBodyRoutePredicateFactory;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.MediaType;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.function.Predicate;

/**
 * 支持空body
 *
 * @author Moonlight
 */

@Component
public class MyReadBodyPredicateFactory extends AbstractRoutePredicateFactory<ReadBodyRoutePredicateFactory.Config> {
  private static final Log log = LogFactory.getLog(MyReadBodyPredicateFactory.class);

  private static final String TEST_ATTRIBUTE = "read_body_predicate_test_attribute";

  public static final String CACHE_REQUEST_BODY_OBJECT_KEY = "cachedRequestBodyObject";

  private final List<HttpMessageReader<?>> messageReaders;

  public MyReadBodyPredicateFactory() {
    super(ReadBodyRoutePredicateFactory.Config.class);
    messageReaders = HandlerStrategies.withDefaults().messageReaders();
  }

  @Override
  @SuppressWarnings("unchecked")
  public AsyncPredicate<ServerWebExchange> applyAsync(final ReadBodyRoutePredicateFactory.Config config) {
    return exchange -> {
      final Class<?> inClass = config.getInClass();
      final Object cachedBody = exchange.getAttribute(CACHE_REQUEST_BODY_OBJECT_KEY);
      if (cachedBody != null) {
        try {
          final boolean test = config.getPredicate().test(cachedBody);
          exchange.getAttributes().put(TEST_ATTRIBUTE, test);
          return Mono.just(test);
        } catch (final ClassCastException e) {
          if (log.isDebugEnabled()) {
            log.debug("Predicate test failed because class in predicate "
              + "does not match the cached body object", e);
          }
        }
        return Mono.just(false);
      } else {
        MediaType mediaType = exchange.getRequest().getHeaders().getContentType();
        if (mediaType != null && !MediaTypeTool.jsonType(mediaType)) {
          return Mono.just(true);
        }
        return ServerWebExchangeUtils.cacheRequestBodyAndRequest(exchange, serverHttpRequest -> ServerRequest
          .create(exchange.mutate().request(serverHttpRequest).build(), messageReaders)
          .bodyToMono(inClass)
          .doOnNext(
            objectValue -> exchange.getAttributes().put(CACHE_REQUEST_BODY_OBJECT_KEY, objectValue))
          .map(objectValue -> config.getPredicate().test(objectValue))).defaultIfEmpty(true);
      }
    };
  }

  @Override
  public Predicate<ServerWebExchange> apply(final ReadBodyRoutePredicateFactory.Config config) {
    throw new UnsupportedOperationException("ReadBodyPredicateFactory is only async.");
  }
}
